#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time    : 2024/6/8 22:11
@Author  : thezehui@gmail.com
@File    : 4.复用提示模板.py
"""
from langchain_core.prompts import PromptTemplate, PipelinePromptTemplate

# 最终模版，需要进行模版内容替换的部分，用变量占位符展位
final_prompt = PromptTemplate.from_template("""
{introduction}

{example}

{start}
""")

introduction_prompt = PromptTemplate.from_template("""
你正在模拟{person}
""")

example_prompt = PromptTemplate.from_template("""
Q: {example_q}
A: {example_a}
""")
start_prompt = PromptTemplate.from_template("""
Q: {input}
A:
""")

pipeline_prompts = [
    ("introduction", introduction_prompt),
    ("example", example_prompt),
    ("start", start_prompt),
]

final = PipelinePromptTemplate(
    final_prompt=final_prompt,
    pipeline_prompts=pipeline_prompts
)

# 第一种方式 这种方式参数是kwargs这种传参采用变量=value的形式，他会自动给把所有的k-v放到一个kwargs参数中
response = final.format(
    person="消毒",
    example_q="中国的首都是哪里？",
    example_a = "北京",
    input = "中国有多少个省份？"
)
# 第二种方式
response_tp = final.format_prompt(
    person="消毒",
    example_q="中国的首都是哪里？",
    example_a = "北京",
    input = "中国有多少个省份？"
)

# 第三种方式，这种参数是dict字段类型
response2 = final.invoke(
    {"person":"消毒",
    "example_q":"中国的首都是哪里？",
    "example_a" : "北京",
    "input" : "中国有多少个省份？"}
)

print(response)
print(response_tp)
print(response_tp.to_string())
print(response2)
